//////////////////////////////////////////////////////////////////////////////////
//	Copyright 2011 by Yosef Grabivker, ryuho@homenetworksoft.com
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.
//////////////////////////////////////////////////////////////////////////////////

#ifdef _WIN32
	#pragma warning (disable:4996)
	#include <Ws2tcpip.h>
	#define SOCK_FLAGS	0 
	#include <mswsock.h>
#else
	#include <netinet/tcp.h>
	#include <netinet/in.h>
	#include <netdb.h>
	#include <errno.h>
	#include <sys/socket.h>
    #include <netinet/in.h>
    #include <arpa/inet.h>
	#include <sys/ioctl.h>
	#include <net/if.h>

	#define SD_BOTH			SHUT_RDWR
	#define SOCKET_ERROR 	-1
	#define SOCK_FLAGS		MSG_NOSIGNAL
	#define WSAETIMEDOUT	EAGAIN 
#endif

#include <stdio.h>


#include "rtsocket.h"

#ifndef INVALID_SOCKET
	#define INVALID_SOCKET -1
#endif


CRTSocket::CRTSocket(int nTimeOut):m_bConnected(false)
{
	int nRes;

	m_socket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);

	if (INVALID_SOCKET==m_socket) THROW_EXC(cERFailed);

	#ifdef LINUX
		timeval iOptVal;
		iOptVal.tv_sec  = nTimeOut/1000;
		iOptVal.tv_usec = (nTimeOut%1000)*1000;
		int iOptLen = sizeof(iOptVal);
	#else
		int iOptVal=nTimeOut; 
		int iOptLen = sizeof(int);
	#endif

	nRes = setsockopt(m_socket, SOL_SOCKET, SO_RCVTIMEO, (const char*)&iOptVal,  iOptLen);
	//printf("%d\n", nRes);
	
	nRes = setsockopt(m_socket, SOL_SOCKET, SO_SNDTIMEO, (const char*)&iOptVal,  iOptLen);
	//printf("%d\n", nRes);
	
}

CRTSocket::~CRTSocket(void)
{
	Close();
}

bool CRTSocket::IsOk()
{
	return m_socket != INVALID_SOCKET;
}

bool CRTSocket::IsConnected()
{
	return m_bConnected;
}

void CRTSocket::Connect(const char* szAddress, unsigned short wPort)
{
	int nConRes=-1;

	if (INVALID_SOCKET==m_socket) THROW_EXC(cERWrongObjectState);

	struct addrinfo *aiList = GetAddress(szAddress, wPort);

	if ( !aiList ) THROW_EXC(cERFailed);

	#if !defined(LINUX) 
	u_long iMode = 1;
	ioctlsocket(m_socket, FIONBIO, &iMode);
	#endif

	nConRes = connect(m_socket, aiList->ai_addr, aiList->ai_addrlen);

	#if !defined(LINUX) 
	if ( nConRes!=0 &&
		 WaitForConnect() )
	{
		nConRes=0;
	}

	iMode=0;
	ioctlsocket(m_socket, FIONBIO, &iMode);
	#endif

	freeaddrinfo( aiList );

	aiList = NULL;

	if (0!=nConRes) THROW_EXC(cERInvalidHostOrPort);

	m_bConnected = true;
}

void CRTSocket::Close()
{
	if (INVALID_SOCKET!=m_socket)
	{
		shutdown(m_socket, SD_BOTH );

		#ifdef _WIN32
			Sleep(500);

			//workaround about close socket win bug http://bugs.sun.com/bugdatabase/view_bug.do?bug_id=4468997
			LPFN_DISCONNECTEX pdisc = NULL;

			DWORD dwBytesReturned; 
			GUID guidDisconnectEx = WSAID_DISCONNECTEX; 
			WSAIoctl( m_socket, 
				SIO_GET_EXTENSION_FUNCTION_POINTER, 
				&guidDisconnectEx, 
				sizeof(GUID), 
				&pdisc, 
				sizeof(pdisc), 
				&dwBytesReturned, 
				NULL, 
				NULL); 

			if (pdisc)
			{
				pdisc(m_socket, NULL, 0, 0); 
			}

			closesocket( m_socket);
		#else
			
			close(m_socket);
		#endif

		m_socket = INVALID_SOCKET;
	}
}


int	 CRTSocket::Write(const char* pData, size_t stSize)
{
	int nRes;

	if (INVALID_SOCKET==m_socket) THROW_EXC(cERWrongObjectState);
	if (!m_bConnected) THROW_EXC(cERWrongObjectState);

	int flags=SOCK_FLAGS;

	nRes = send(m_socket, pData, (int)stSize, flags);

	if ( SOCKET_ERROR==nRes )
	{
		CheckError();
	}

	return nRes;
}

int  CRTSocket::Read( char* pData, size_t stSize)
{
	int nRes;

	if (INVALID_SOCKET==m_socket) THROW_EXC(cERWrongObjectState);
	if (!m_bConnected) THROW_EXC(cERWrongObjectState);

	int flags=SOCK_FLAGS;

	nRes = recv(m_socket, (char*)pData, (int)stSize, flags);

	if ( SOCKET_ERROR==nRes )
	{
		CheckError();

		nRes = 0;
	}

	return nRes;
}

bool CRTSocket::WaitForRead(long milliseconds )
{
	if (INVALID_SOCKET==m_socket) THROW_EXC(cERWrongObjectState);
	if (!m_bConnected) THROW_EXC(cERWrongObjectState);

	bool bRes(false);
	int res;
	fd_set          fdRead; 
	struct timeval  timeout;

	FD_ZERO(&fdRead); 
	FD_SET((unsigned int)m_socket, &fdRead); 

	memset((char *)&timeout,0,sizeof(timeout));

	timeout.tv_sec  = milliseconds/1000;
	timeout.tv_usec = (milliseconds%1000)*1000;

	//if res==0 timeout
	res = select(m_socket+1, &fdRead, NULL, NULL, &timeout); 

	if( FD_ISSET(m_socket, &fdRead) ) 
		bRes = true;
	else
	{
		if (SOCKET_ERROR==res) CheckError();
	}
	
	return bRes;
}

bool CRTSocket::WaitForConnect(long milliseconds)
{
	if (INVALID_SOCKET==m_socket) THROW_EXC(cERWrongObjectState);

	bool bRes(false);
	int res;
	fd_set          fdRead; 
	struct timeval  timeout;

	FD_ZERO(&fdRead); 
	FD_SET((unsigned int)m_socket, &fdRead); 

	memset((char *)&timeout,0,sizeof(timeout));

	timeout.tv_sec  = milliseconds/1000;
	timeout.tv_usec = (milliseconds%1000)*1000;

	//if res==0 timeout
	res = select(m_socket+1, &fdRead, NULL, NULL, &timeout); 

	if( FD_ISSET(m_socket, &fdRead) ) 
		bRes = true;
	
	return bRes;
}

void CRTSocket::CheckError()
{
#ifdef _WIN32
	int nErr = WSAGetLastError();
#else
	int nErr = errno;
	
	switch ( nErr )
	{
		case  	EAGAIN:
		{
			return;
		}
		break;
	}
	
	
#endif
	
	if ( nErr && nErr!=WSAETIMEDOUT  )
	{
		m_bConnected = false;
		THROW_EXC(cERWrongObjectState);
	}
}


void CRTSocket::Initialize()
{
#ifdef _WIN32
	WORD wVersionRequested;
	WSADATA wsaData;
	int err;

	wVersionRequested = MAKEWORD( 2, 0 );

	err = WSAStartup( wVersionRequested, &wsaData );

	if ( err != 0 )
	{
		exit(1);
	}
#endif
}

struct addrinfo* CRTSocket::GetAddress( const char* szAddr, unsigned short wPort )
{
	char buf[32];
	sprintf(buf, "%u", (unsigned int)wPort);

	struct addrinfo aiHints;
	struct addrinfo *aiList = NULL;
	memset(&aiHints, 0, sizeof(aiHints));
	aiHints.ai_family = AF_INET;
	aiHints.ai_socktype = SOCK_STREAM;
	aiHints.ai_protocol = IPPROTO_TCP;

	if ( 0==getaddrinfo(szAddr, buf, &aiHints, &aiList) )
	{
		return aiList;
	}

	return NULL;
}

bool CRTSocket::IsAddressEqual( const char* szAddr1, const char* szAddr2 )
{
	bool bRes(false);

	struct addrinfo* a = CRTSocket::GetAddress( szAddr1 );
	struct addrinfo* b = CRTSocket::GetAddress( szAddr2 );

	if ( a && b )
	{
		if ( 0==memcmp(a->ai_addr, b->ai_addr, a->ai_addrlen) )
		{
			bRes = true;
		}
	}

	if (a) freeaddrinfo( a );
	if (b) freeaddrinfo( b );

	return bRes;
}

void CRTSocket::Purge()
{
	char buf[1024];

	//must be while here
	while ( WaitForRead(25) )
	{
		if ( Read(buf, sizeof(buf))<=0 )
			break;
	}
}
